import os
import argparse

def convert_txt_to_xml(source_txt_path, source_img_path, save_xml_path):
    # 创建保存XML文件的目录
    if not os.path.exists(save_xml_path):
        os.makedirs(save_xml_path)
    
    # 读取类别名称，假设是通过文件名分隔符获得的
    def get_class_name(file_name):
        return file_name.split("_2023")[0]

    # 遍历文本文件，生成对应的XML文件
    for txt_file in os.listdir(source_txt_path):
        if txt_file.endswith(".txt"):
            with open(os.path.join(source_txt_path, txt_file), "r") as txt:
                lines = txt.readlines()

            xml_file_name = os.path.splitext(txt_file)[0] + ".xml"
            xml_path = os.path.join(save_xml_path, xml_file_name)
            
            with open(xml_path, "w") as xml:
                xml.write('<?xml version="1.0" encoding="UTF-8"?>\n')
                xml.write('<annotations>\n')

                for line in lines:
                    line_split = line.strip().split()
                    img_name = line_split[0]
                    img_path = os.path.join(source_img_path, img_name + ".jpg")
                    class_name = get_class_name(img_name)
                    
                    xml.write('\t<object>\n')
                    xml.write('\t\t<name>' + class_name + '</name>\n')
                    xml.write('\t\t<image_path>' + img_path + '</image_path>\n')
                    xml.write('\t\t<bounding_box>' + ' '.join(line_split[1:]) + '</bounding_box>\n')
                    xml.write('\t</object>\n')

                xml.write('</annotations>\n')

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Convert txt annotations to xml format")
    parser.add_argument("--source_txt_path", type=str, required=True, help="Path to source txt files")
    parser.add_argument("--source_img_path", type=str, required=True, help="Path to source image files")
    parser.add_argument("--save_xml_path", type=str, required=True, help="Path to save xml files")
    args = parser.parse_args()

    convert_txt_to_xml(args.source_txt_path, args.source_img_path, args.save_xml_path)

# python txt_to_xml.py --source_txt_path 20230618_wangjian_200_labels --source_img_path 20230618_wangjian_200_jpg --save_xml_path 20230618_wangjian_200_xml
